package com.plugin.dao.config;

import java.sql.ResultSet;
import java.util.List;
import java.util.Set;

import javax.sql.DataSource;

import org.apache.shardingsphere.core.rule.TableRule;
import org.apache.shardingsphere.shardingjdbc.jdbc.core.connection.ShardingConnection;
import org.apache.shardingsphere.shardingjdbc.jdbc.core.context.ShardingRuntimeContext;
import org.apache.shardingsphere.shardingjdbc.jdbc.core.datasource.ShardingDataSource;
import org.apache.shardingsphere.shardingjdbc.spring.boot.SpringBootConfiguration;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.autoconfigure.AutoConfigureAfter;
import org.springframework.context.EnvironmentAware;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.env.Environment;

import com.google.common.collect.Lists;

@Configuration
@AutoConfigureAfter(SpringBootConfiguration.class)
public class ShardingProperty implements EnvironmentAware {

    private final static String GET_TABLE_SQL = "select table_name from information_schema.tables where table_schema='%s'";
    private final static String TABLE_NAME_COLUMN = "table_name";

    @Autowired
    private DataSource dataSource;


    @Override
    public void setEnvironment(Environment environment) {
        ShardingDataSource shardingDataSource = (ShardingDataSource) dataSource;
        ShardingRuntimeContext runtimeContext = shardingDataSource.getRuntimeContext();
        Set<String> dsSet = shardingDataSource.getDataSourceMap().keySet();
        for (String ds : dsSet) {
            String schema = runtimeContext.getMetaData().getDataSources().getDataSourceMetaData(ds).getCatalog();
            ShardingConnection connection = shardingDataSource.getConnection();
            try {
                ResultSet resultSet = connection.createStatement().executeQuery(String.format(GET_TABLE_SQL, schema));
                List<String> tableNameList = Lists.newArrayList();
                while (resultSet.next()) {
                    String tableName = resultSet.getString(TABLE_NAME_COLUMN);
                    tableNameList.add(tableName);
                }
                for (String tableName : tableNameList) {
                    if (!shardingDataSource.getRuntimeContext().getRule().findTableRuleByActualTable(tableName).isPresent()) {
                        TableRule rule = new TableRule(dsSet, tableName);
                        shardingDataSource.getRuntimeContext().getRule().getTableRules().add(rule);
                    }
                }
            }
            catch (Exception e) {
                e.printStackTrace();
            }
            finally {
                try {
                    if (!connection.isClosed()) {
                        connection.close();
                    }
                }
                catch (Exception e) {
                    e.printStackTrace();
                }
            }
        }
    }
}
